#!/usr/bin/env python
# encoding: utf-8

import torch
import torch.nn as nn
import torch.nn.functional as F

class TDNN(nn.Module):
    def __init__(self, nOut, **kwargs):
        super(TDNN, self).__init__()
        self.td_layer1 = torch.nn.Conv1d(in_channels=64, out_channels=256, dilation=1, kernel_size=5, stride=1)
        self.bn1 = nn.BatchNorm1d(256)
        self.td_layer2 = torch.nn.Conv1d(in_channels=256, out_channels=512, dilation=2, kernel_size=3, stride=1)
        self.bn2 = nn.BatchNorm1d(512)
        self.td_layer3 = torch.nn.Conv1d(in_channels=512, out_channels=512, dilation=3, kernel_size=3, stride=1)
        self.bn3 = nn.BatchNorm1d(512)
        self.td_layer4 = torch.nn.Conv1d(in_channels=512, out_channels=512, dilation=1, kernel_size=1, stride=1)
        self.bn4 = nn.BatchNorm1d(512)
        self.td_layer5 = torch.nn.Conv1d(in_channels=512, out_channels=1500, dilation=1, kernel_size=1, stride=1)
        self.bn5 = nn.BatchNorm1d(1500)

        self.fc1 = nn.Linear(3000, 512)
        self.fc2 = nn.Linear(512, nOut)

    def forward(self, x):
        '''
        x [batch_size, dim, time]
        '''
        x = F.relu(self.td_layer1(x))
        x = self.bn1(x)

        x = F.relu(self.td_layer2(x))
        x = self.bn2(x)

        x = F.relu(self.td_layer3(x))
        x = self.bn3(x)

        x = F.relu(self.td_layer4(x))
        x = self.bn4(x)

        x = F.relu(self.td_layer5(x))
        x = self.bn5(x)

        x = self.statistics_pooling(x)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)

        return x

    @staticmethod
    def statistics_pooling(x):
        """Computes Statistics Pooling Module
        Args:
            x (torch.Tensor): Input tensor (#batch, channels, time).
        Returns:
            torch.Tensor: Output tensor (#batch, channels*2)
        """
        mean = torch.mean(x, axis=2)
        var = torch.var(x, axis=2)
        x = torch.cat((mean, var), axis=1)
        return x


def Speaker_Encoder(nOut=256, **kwargs):
    model = TDNN(nOut, **kwargs)
    return model


